
# from PIL import Image
from pymage_size import get_image_size

import glob
import logging
import os

import helper


logger = logging.getLogger(__name__)


class Main():

    def __init__(self, args):
        self.args = args

    def start(self):
        class_names = []

        os.makedirs(self.args.out_dir, exist_ok=True)

        if self.args.classes_remap:
            remap_filename, varname = self.args.classes_remap.rsplit(':', maxsplit=1)

            remap_file = helper.dynamic_import('remap_file', remap_filename)
            classes_remap = getattr(remap_file, varname)
        else:
            classes_remap = None

        print('classes_remap', classes_remap)
        # return

        for index, label_fullname in enumerate(glob.glob(os.path.join(self.args.lbl_dir, '*.txt'))):
            basename = os.path.basename(label_fullname)

            img_fullname = os.path.join(self.args.img_dir, basename.replace('.txt', self.args.img_ext))
            if not os.path.exists(img_fullname):
                logger.warning('image not found: %s', img_fullname)
                continue

            # img = Image.open(img_fullname)
            # iw, ih = img.size
            img = get_image_size(img_fullname)
            iw, ih = img.get_dimensions()
            img_hw = (ih, iw)

            parser = helper.DotaObbTxt(class_names, classes_remap, label_fullname)

            boxes = []
            for class_id, cxywh in parser.get_boxes(img_hw):
                nxywh = helper.normalize_xywh(cxywh, img_hw)
                boxes.append((class_id, nxywh))

            if args.skip_empty and not boxes:
                continue

            out_fullname = os.path.join(self.args.out_dir, basename)
            with open(out_fullname, 'w') as f:
                for class_id, nxywh in boxes:
                    f.write('%d %f %f %f %f\n' % (class_id, *nxywh))

            if args.debug:
                print(basename, len(boxes))
                break
        
            if index % 100 == 0:
                print(index, basename, len(boxes))

        if self.args.classes_file is not None:
            classes_fullname = self.args.classes_file
        else:
            classes_fullname = os.path.join(self.args.out_dir, 'classes.txt')
        with open(classes_fullname, 'w') as f:
            for class_name in class_names:
                f.write('%s\n' % class_name)


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    # trainsplit
    # valsplit
    default_root = r'D:\static\datasets\DOTA_split\\valsplit'
    parser.add_argument('--img-dir', type=str, default=f'{default_root}\\images')
    parser.add_argument('--img-ext', type=str, default='.jpg')
    parser.add_argument('--lbl-dir', type=str, default=f'{default_root}\\labelTxt')
    parser.add_argument('--out-dir', type=str, default=f'{default_root}\\labels')
    parser.add_argument('--classes-remap', type=str, default=r'.\labelfmt\remap_data.py:DOTA_split_to_r3is')
    parser.add_argument('--classes-file', type=str, default=r'{default_root}\\classes.txt')
    parser.add_argument('--debug', action='store_true')
    parser.add_argument('--skip-empty', action='store_true')
    args = parser.parse_args()
    Main(args).start()
